package com.chenfan.mcn.config;

import com.chenfan.infra.context.TenantContextHolder;
import com.chenfan.common.extension.context.TraceContext;
import com.chenfan.common.extension.context.TraceContextHolder;
import com.chenfan.common.extension.util.SmartExecutorService;
import lombok.extern.slf4j.Slf4j;
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.task.TaskDecorator;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;

import java.util.concurrent.Executor;
import java.util.concurrent.ExecutorService;

/**
 * @author: tangwei
 * @Date: 2021/7/21 15:35
 * @Description:
 */
@Slf4j
@Configuration
public class ThreadPoolConfig {


    @Bean("threadPoolTaskExecutor")
    public Executor asyncExecutor() {
        ThreadPoolTaskExecutor threadPoolTaskExecutor = new ThreadPoolTaskExecutor();
        threadPoolTaskExecutor.setCorePoolSize(PoolProperties.CORE_SIZE);
        threadPoolTaskExecutor.setMaxPoolSize(PoolProperties.MAX_SIZE);
        threadPoolTaskExecutor.setQueueCapacity(PoolProperties.QUEUE_CAPACITY);
        threadPoolTaskExecutor.setThreadNamePrefix(PoolProperties.THREAD_NAME_PREFIX);
        threadPoolTaskExecutor.setTaskDecorator(new AsyncTaskDecorator());
        threadPoolTaskExecutor.initialize();
        return threadPoolTaskExecutor;
    }

    class PoolProperties {

        private static final int CORE_SIZE = 1;

        private static final int MAX_SIZE = 16;

        private static final int QUEUE_CAPACITY = 100;

        private static final String THREAD_NAME_PREFIX = "mcn-async-thread-";

    }

    class AsyncTaskDecorator implements TaskDecorator {

        @Override
        public Runnable decorate(Runnable runnable) {
            TraceContext traceContext = TraceContextHolder.getCurrentContext();
            Long tenantId = UserVoContextHolder.getTenantId();
            return () -> {
                try {
                    log.info("=== TaskDecorator tenantId: {} ===", tenantId);
                    TraceContextHolder.setContext(traceContext);
                    TenantContextHolder.setTenantId(tenantId);
                    runnable.run();
                } catch (Exception e) {
                    log.error("AsyncTaskDecorator.decorate error:{}", e);
                } finally {
                    TraceContextHolder.clear();
                    TenantContextHolder.clear();
                }

            };
        }
    }

    @Bean(name = "executorService")
    @ConditionalOnMissingBean(ExecutorService.class)
    public ExecutorService executorService() {
        return SmartExecutorService.newFixedThreadPool();
    }
}